import sys
import os
if 'win' in sys.platform:
    bert_model_path = r"F:\pretrain_model\bert\chinese_L-12_H-768_A-12"
else:
    bert_model_path = r'/data/nlp_data/bert/model_chinese/'
import sys
sys.path.append("/data/python_project/qiufengfeng/nlp_tools")
from nlp_tools.third_model.bert4keras.bert4keras.models import build_transformer_model
from nlp_tools.third_model.bert4keras.bert4keras.tokenizers import Tokenizer
import numpy as np


config_path = os.path.join(bert_model_path,'bert_config.json')
check_point_path = os.path.join(bert_model_path,'bert_model.ckpt')
vocab_path = os.path.join(bert_model_path,'vocab.txt')

# 建立分词器
tokenizer = Tokenizer(vocab_path,do_lower_case=True)
model = build_transformer_model(config_path=config_path,
                                checkpoint_path=check_point_path,
                                with_mlm=True)
token_ids, segment_ids = tokenizer.encode('科学技术是第一生产力')
token_ids[3] = token_ids[4] = tokenizer._token_dict['[MASK]']

# 用mlm模型预测被mask掉的部分
probas = model.predict([np.array([token_ids]), np.array([segment_ids])])[0]
print(tokenizer.decode(probas[3:5].argmax(axis=1)))  # 结果正是“技术”